import tkinter as tk
from tkinter import filedialog, simpledialog, messagebox
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import os

# 水印嵌入函数
def embed_watermark(image_path, watermark, output_path, position=(0, 0), alpha=128):
    try:
        base_image = Image.open(image_path).convert("RGBA")
        watermark_image = None

        if isinstance(watermark, str):  # 如果水印是文本
            watermark_image = Image.new("RGBA", base_image.size, (255, 255, 255, 0))
            draw = ImageDraw.Draw(watermark_image)
            font_path = get_font_path()  # 确保字体文件存在
            if not os.path.exists(font_path):
                raise FileNotFoundError(f"Font file {font_path} not found.")
            font = ImageFont.truetype(font_path, size=40)

            # 使用 font.getbbox 属性来获取文本尺寸
            bbox = font.getbbox(watermark)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]

            # 确保水印不会超出图像边界
            position = (max(position[0], 0), max(position[1], 0))
            draw.text((position[0], position[1]), watermark, fill=(255, 255, 255, alpha), font=font)
        else:  # 如果水印是图像
            watermark_image = Image.open(watermark).convert("RGBA")
            watermark_image = watermark_image.resize((base_image.width // 10, base_image.height // 10))
            watermark_image.putalpha(alpha)

        watermarked_image = Image.alpha_composite(base_image, watermark_image)
        watermarked_image.convert("RGB").save(output_path)
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred: {e}")

# 获取字体路径的函数
def get_font_path():
    # 尝试使用默认的 Arial 字体路径
    default_font_paths = [
        "arial.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "C:\\Windows\\Fonts\\arial.ttf"
    ]
    for path in default_font_paths:
        if os.path.exists(path):
            return path

    # 如果默认路径找不到，让用户选择字体文件
    font_path = filedialog.askopenfilename(title="Select a font file", filetypes=[("Font files", "*.ttf")])
    return font_path

# 水印提取函数
def extract_watermark(image_path, watermark_output_path, position, watermark_size):
    watermarked_image = Image.open(image_path).convert("RGBA")
    watermark_region = watermarked_image.crop((position[0], position[1], position[0] + watermark_size[0], position[1] + watermark_size[1]))
    watermark_region.save(watermark_output_path)

# 创建GUI
class WatermarkApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Watermark Embedder")

        # 创建按钮和文本框
        self.original_image_label = tk.Label(root, text="Original Image:")
        self.original_image_label.grid(row=0, column=0)
        self.original_image_entry = tk.Entry(root)
        self.original_image_entry.grid(row=0, column=1)
        self.original_image_button = tk.Button(root, text="Open", command=self.open_original_image)
        self.original_image_button.grid(row=0, column=2)

        self.watermark_label = tk.Label(root, text="Watermark (Text/Image):")
        self.watermark_label.grid(row=1, column=0)
        self.watermark_entry = tk.Entry(root)
        self.watermark_entry.grid(row=1, column=1)
        self.watermark_button = tk.Button(root, text="Set", command=self.set_watermark)
        self.watermark_button.grid(row=1, column=2)

        self.embed_button = tk.Button(root, text="Embed Watermark", command=self.embed_watermark)
        self.embed_button.grid(row=2, column=1)

    def open_original_image(self):
        self.original_image_path = filedialog.askopenfilename()
        self.original_image_entry.delete(0, tk.END)
        self.original_image_entry.insert(0, self.original_image_path)

    def set_watermark(self):
        self.watermark = simpledialog.askstring("Input", "Enter watermark text or select an image:", parent=self.root)
        if self.watermark:
            self.watermark_entry.delete(0, tk.END)
            self.watermark_entry.insert(0, self.watermark)
        else:
            self.watermark = filedialog.askopenfilename()
            self.watermark_entry.delete(0, tk.END)
            self.watermark_entry.insert(0, self.watermark)

    def embed_watermark(self):
        if self.original_image_path and self.watermark:
            output_path = filedialog.asksaveasfilename(defaultextension=".jpg")
            if output_path:
                # 允许用户选择水印位置
                position = simpledialog.askstring("Input", "Enter watermark position (x,y):", parent=self.root)
                if position:
                    position = tuple(map(int, position.split(',')))
                else:
                    position = (0, 0)

                embed_watermark(self.original_image_path, self.watermark, output_path, position=position)
                messagebox.showinfo("Success", "Watermark embedded successfully!")
        else:
            messagebox.showerror("Error", "Please select an image and set a watermark.")

# 主函数
def main():
    root = tk.Tk()
    app = WatermarkApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()
